{
 "metadata": {
  "name": "",
  "signature": "sha256:efebca98f3b61822c046b6fd4aa23c699b738e515f5ed666d4a6afc3e90b2dd3"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# visualization of a single qubit state\n",
      "# James Marshall\n",
      "\n",
      "import svgwrite\n",
      "import math\n",
      "\n",
      "from IPython.display import display, HTML, clear_output\n",
      "from IPython.html.widgets import FloatSlider, HTMLWidget, ContainerWidget, Latex, Checkbox\n",
      "\n",
      "# state variables\n",
      "a1, b1, rho1, theta1 = 1, 0, 1, 0\n",
      "a2, b2, rho2, theta2 = 0, 0, 0, 0\n",
      "last_sign_a1, last_sign_b1, last_sign_a2, last_sign_b2 = 1, 1, 1, 1\n",
      "\n",
      "def spaces(n):\n",
      "    return HTMLWidget(value=r\" \"*n)\n",
      "\n",
      "def qubit_latex_expression():\n",
      "    phase1 = \"\" if theta1 == 0 else \"e^{%g i}\" % round(theta1, 2)\n",
      "    phase2 = \"\" if theta2 == 0 else \"e^{%g i}\" % round(theta2, 2)\n",
      "    return \"$\\\\vert\\\\psi\\\\rangle = %g %s \\\\vert 0 \\\\rangle + %g %s\\\\vert 1 \\\\rangle$\" \\\n",
      "        % (round(rho1, 2), phase1, round(rho2, 2), phase2)\n",
      "\n",
      "callbacksDisabled = False\n",
      "\n",
      "# create sliders\n",
      "a1_slider = FloatSliderWidget(min=-1, max=1, value=a1, description=\"a1\", step=0.01)\n",
      "b1_slider = FloatSliderWidget(min=-1, max=1, value=b1, description=\"b1\", orientation=\"horizontal\", step=0.01)\n",
      "rho1_slider = FloatSliderWidget(min=0, max=1, value=rho1, description=\"magnitude1\", step=0.01)\n",
      "theta1_slider = FloatSliderWidget(min=0, max=round(2*math.pi, 2), value=theta1, description=\"phase1\", step=0.01)\n",
      "\n",
      "a2_slider = FloatSliderWidget(min=-1, max=1, value=a2, description=\"a2\", step=0.01)\n",
      "b2_slider = FloatSliderWidget(min=-1, max=1, value=b2, description=\"b2\", step=0.01)\n",
      "rho2_slider = FloatSliderWidget(min=0, max=1, value=rho2, description=\"magnitude2\", step=0.01)\n",
      "theta2_slider = FloatSliderWidget(min=0, max=round(2*math.pi, 2), value=theta2, description=\"phase2\", step=0.01)\n",
      "\n",
      "prob_slider = FloatSliderWidget(min=0, max=1, value=0, description=\"Probability of |1>\", step=0.01)\n",
      "latitude_slider = FloatSliderWidget(min=0, max=round(math.pi, 2), value=0, description=\"Bloch sphere latitude (theta)\", step=0.01)\n",
      "longitude_slider = FloatSliderWidget(min=0, max=round(2*math.pi, 2), value=0, description=\"Bloch sphere longitude (phi)\", step=0.01)\n",
      "\n",
      "c1_label = LatexWidget(value=r\"$c_1 = a_1 + b_1 i = \\rho_1 e^{i\\theta_1}$\")\n",
      "c2_label = LatexWidget(value=r\"$c_2 = a_2 + b_2 i = \\rho_2 e^{i\\theta_2}$\")\n",
      "\n",
      "qubit_label = LatexWidget(value=qubit_latex_expression())\n",
      "\n",
      "c1_container = ContainerWidget(children=[c1_label, rho1_slider, theta1_slider, a1_slider, b1_slider])\n",
      "c2_container = ContainerWidget(children=[c2_label, rho2_slider, theta2_slider, a2_slider, b2_slider])\n",
      "\n",
      "top_container = ContainerWidget(children=[c1_container, spaces(20), c2_container])\n",
      "\n",
      "image = HTMLWidget()\n",
      "image_container = ContainerWidget(children=[spaces(60), image])\n",
      "\n",
      "prob_container = ContainerWidget(children=[spaces(60), prob_slider])\n",
      "\n",
      "latitude_container = ContainerWidget(children=[spaces(42), latitude_slider])\n",
      "longitude_container = ContainerWidget(children=[spaces(42), longitude_slider])\n",
      "\n",
      "bloch_container = ContainerWidget(children=[spaces(1), latitude_container, longitude_container])\n",
      "\n",
      "qubit_container = ContainerWidget(children=[spaces(100), qubit_label])\n",
      "\n",
      "phases_coupled = CheckboxWidget(value=False, description=\"Phases coupled\")\n",
      "\n",
      "all = ContainerWidget(children=[top_container, phases_coupled, qubit_container, image_container, prob_container, bloch_container])\n",
      "display(all)\n",
      "\n",
      "# top_container.remove_class('vbox')\n",
      "# top_container.add_class('hbox')\n",
      "\n",
      "# qubit_container.remove_class('vbox')\n",
      "# qubit_container.add_class('hbox')\n",
      "\n",
      "# image_container.remove_class('vbox')\n",
      "# image_container.add_class('hbox')\n",
      "\n",
      "# prob_container.remove_class('vbox')\n",
      "# prob_container.add_class('hbox')\n",
      "\n",
      "# latitude_container.remove_class('vbox')\n",
      "# latitude_container.add_class('hbox')\n",
      "# longitude_container.remove_class('vbox')\n",
      "# longitude_container.add_class('hbox')\n",
      "\n",
      "\n",
      "def theta(a, b):\n",
      "    if a > 0 and b == 0:\n",
      "        return 0\n",
      "    elif a == 0 and b > 0:\n",
      "        return math.pi/2\n",
      "    elif a < 0 and b == 0:\n",
      "        return math.pi\n",
      "    elif a == 0 and b < 0:\n",
      "        return 3*math.pi/2\n",
      "    elif a > 0 and b > 0:\n",
      "        return math.atan(float(b)/a)\n",
      "    elif a < 0 and b > 0:\n",
      "        return math.atan(float(b)/a) + math.pi\n",
      "    elif a < 0 and b < 0:\n",
      "        return math.atan(float(b)/a) + math.pi\n",
      "    elif a > 0 and b < 0:\n",
      "        return math.atan(float(b)/a) + 2*math.pi\n",
      "    else:\n",
      "        return None\n",
      "    \n",
      "def modulo2pi(theta):\n",
      "    if theta < 0:\n",
      "        theta += 2*math.pi\n",
      "    elif theta >= 2*math.pi:\n",
      "        theta -= 2*math.pi\n",
      "    return theta\n",
      "\n",
      "def rectangular(rho, theta):\n",
      "    a = rho * math.cos(theta)\n",
      "    b = rho * math.sin(theta)\n",
      "    return a, b\n",
      "\n",
      "def update_a1(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    a1 = min(max(new_trait_value, -rho1), rho1)\n",
      "    b1 = math.sqrt(rho1**2 - a1**2) * last_sign_b1\n",
      "    if rho1 > 0:\n",
      "        theta1, previous = theta(a1, b1), theta1\n",
      "        if phases_coupled.value == True:\n",
      "            delta = theta1 - previous\n",
      "            theta2 = modulo2pi(theta2 + delta)\n",
      "            a2, b2 = rectangular(rho2, theta2)\n",
      "    redraw()\n",
      "\n",
      "def update_b1(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    b1 = min(max(new_trait_value, -rho1), rho1)\n",
      "    a1 = math.sqrt(rho1**2 - b1**2) * last_sign_a1\n",
      "    if rho1 > 0:\n",
      "        theta1, previous = theta(a1, b1), theta1\n",
      "        if phases_coupled.value == True:\n",
      "            delta = theta1 - previous\n",
      "            theta2 = modulo2pi(theta2 + delta)\n",
      "            a2, b2 = rectangular(rho2, theta2)\n",
      "    redraw()\n",
      "    \n",
      "def update_rho1(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    rho1 = new_trait_value\n",
      "    a1, b1 = rectangular(rho1, theta1)\n",
      "    rho2 = math.sqrt(1 - rho1**2)\n",
      "    a2, b2 = rectangular(rho2, theta2)\n",
      "    redraw()\n",
      "\n",
      "def update_theta1(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    theta1, previous = new_trait_value, theta1\n",
      "    a1, b1 = rectangular(rho1, theta1)\n",
      "    if phases_coupled.value == True:\n",
      "        delta = theta1 - previous\n",
      "        theta2 = modulo2pi(theta2 + delta)\n",
      "        a2, b2 = rectangular(rho2, theta2)\n",
      "    redraw()\n",
      "\n",
      "def update_a2(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    a2 = min(max(new_trait_value, -rho2), rho2)\n",
      "    b2 = math.sqrt(rho2**2 - a2**2) * last_sign_b2\n",
      "    if rho2 > 0:\n",
      "        theta2, previous = theta(a2, b2), theta2\n",
      "        if phases_coupled.value == True:\n",
      "            delta = theta2 - previous\n",
      "            theta1 = modulo2pi(theta1 + delta)\n",
      "            a1, b1 = rectangular(rho1, theta1)\n",
      "    redraw()\n",
      "\n",
      "def update_b2(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    b2 = min(max(new_trait_value, -rho2), rho2)\n",
      "    a2 = math.sqrt(rho2**2 - b2**2) * last_sign_a2\n",
      "    if rho2 > 0:\n",
      "        theta2, previous = theta(a2, b2), theta2\n",
      "        if phases_coupled.value == True:\n",
      "            delta = theta2 - previous\n",
      "            theta1 = modulo2pi(theta1 + delta)\n",
      "            a1, b1 = rectangular(rho1, theta1)\n",
      "    redraw()\n",
      "    \n",
      "def update_rho2(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    rho2 = new_trait_value\n",
      "    a2, b2 = rectangular(rho2, theta2)\n",
      "    rho1 = math.sqrt(1 - rho2**2)\n",
      "    a1, b1 = rectangular(rho1, theta1)\n",
      "    redraw()\n",
      "\n",
      "def update_theta2(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    theta2, previous = new_trait_value, theta2\n",
      "    a2, b2 = rectangular(rho2, theta2)\n",
      "    if phases_coupled.value == True:\n",
      "        delta = theta2 - previous\n",
      "        theta1 = modulo2pi(theta1 + delta)\n",
      "        a1, b1 = rectangular(rho1, theta1)\n",
      "    redraw()\n",
      "\n",
      "def update_prob(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    rho2 = math.sqrt(new_trait_value)\n",
      "    a2, b2 = rectangular(rho2, theta2)\n",
      "    rho1 = math.sqrt(1 - rho2**2)\n",
      "    a1, b1 = rectangular(rho1, theta1)\n",
      "    redraw()\n",
      "\n",
      "def update_latitude(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    latitude = new_trait_value\n",
      "    rho1 = math.cos(latitude/2.0)\n",
      "    rho2 = math.sin(latitude/2.0)\n",
      "    a1, b1 = rectangular(rho1, theta1)\n",
      "    a2, b2 = rectangular(rho2, theta2)\n",
      "    redraw()\n",
      "\n",
      "def update_longitude(trait_name, new_trait_value):\n",
      "    if callbacksDisabled: return\n",
      "    global a1, b1, rho1, theta1, a2, b2, rho2, theta2\n",
      "    if phases_coupled.value == False:\n",
      "        longitude = new_trait_value\n",
      "        theta2 = modulo2pi(theta1 + longitude)\n",
      "        a2, b2 = rectangular(rho2, theta2)\n",
      "    redraw()\n",
      "\n",
      "\n",
      "def redraw():\n",
      "    global callbacksDisabled, last_sign_a1, last_sign_b1, last_sign_a2, last_sign_b2\n",
      "    clear_output(wait=True)\n",
      "    d = svgwrite.Drawing(size=(600, 300), debug=False)\n",
      "    d.viewbox(-1.25, -1.25, 6, 2.5)    # xmin, ymin, width, height\n",
      "    # c1\n",
      "    d.add(d.circle(center=(0,0), r=1, fill='white', stroke='black', stroke_width=0.01))\n",
      "    # arrow\n",
      "    d.add(d.line(start=(0,0), end=(a1, -b1), stroke='blue', stroke_width=0.03))\n",
      "    # arrowhead\n",
      "    d.add(d.circle(center=(a1, -b1), r=0.03, fill='red'))\n",
      "    # label\n",
      "    d.add(d.text(\"c1\", insert=(-0.1,-1.1), font_size=0.25))\n",
      "    # projections\n",
      "    d.add(d.line(start=(0,0), end=(a1, 0), stroke='red', stroke_width=0.01))\n",
      "    d.add(d.line(start=(0,0), end=(0, -b1), stroke='red', stroke_width=0.01))\n",
      "\n",
      "    # c2\n",
      "    off = 3   # c2 offset\n",
      "    d.add(d.circle(center=(0+off,0), r=1, fill='white', stroke='black', stroke_width=0.01))\n",
      "    # arrow\n",
      "    d.add(d.line(start=(0+off,0), end=(a2+off, -b2), stroke='blue', stroke_width=0.03))\n",
      "    # arrowhead\n",
      "    d.add(d.circle(center=(a2+off, -b2), r=0.03, fill='red'))\n",
      "    # label\n",
      "    d.add(d.text(\"c2\", insert=(-0.1+off,-1.1), font_size=0.25))\n",
      "    d.add(d.line(start=(0+off,0), end=(a2+off, 0), stroke='red', stroke_width=0.01))\n",
      "    d.add(d.line(start=(0+off,0), end=(0+off, -b2), stroke='red', stroke_width=0.01))\n",
      "\n",
      "    # update signs if non-zero\n",
      "    last_sign_a1 = +1 if a1 > 0 else -1 if a1 < 0 else last_sign_a1\n",
      "    last_sign_b1 = +1 if b1 > 0 else -1 if b1 < 0 else last_sign_b1\n",
      "    last_sign_a2 = +1 if a2 > 0 else -1 if a2 < 0 else last_sign_a2\n",
      "    last_sign_b2 = +1 if b2 > 0 else -1 if b2 < 0 else last_sign_b2\n",
      "\n",
      "    # update qubit expression\n",
      "    qubit_label.value = qubit_latex_expression()\n",
      "\n",
      "    # update sliders\n",
      "    callbacksDisabled = True\n",
      "    a1_slider.value = round(a1, 2)\n",
      "    b1_slider.value = round(b1, 2)\n",
      "    rho1_slider.value = round(rho1, 2)\n",
      "    theta1_slider.value = round(theta1, 2)\n",
      "    a2_slider.value = round(a2, 2)\n",
      "    b2_slider.value = round(b2, 2)\n",
      "    rho2_slider.value = round(rho2, 2)\n",
      "    theta2_slider.value = round(theta2, 2)\n",
      "    prob_slider.value = round(rho2*rho2, 2)\n",
      "    latitude_slider.value = round(2*math.acos(rho1), 2)\n",
      "    longitude_slider.value = round(modulo2pi(theta2-theta1), 2)\n",
      "    callbacksDisabled = False\n",
      "\n",
      "    image.value = d.tostring()\n",
      "    \n",
      "\n",
      "a1_slider.on_trait_change(update_a1, 'value')\n",
      "b1_slider.on_trait_change(update_b1, 'value')\n",
      "rho1_slider.on_trait_change(update_rho1, 'value')\n",
      "theta1_slider.on_trait_change(update_theta1, 'value')\n",
      "a2_slider.on_trait_change(update_a2, 'value')\n",
      "b2_slider.on_trait_change(update_b2, 'value')\n",
      "rho2_slider.on_trait_change(update_rho2, 'value')\n",
      "theta2_slider.on_trait_change(update_theta2, 'value')\n",
      "prob_slider.on_trait_change(update_prob, 'value')\n",
      "latitude_slider.on_trait_change(update_latitude, 'value')\n",
      "longitude_slider.on_trait_change(update_longitude, 'value')\n",
      "redraw()\n",
      "    \n"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 10
    }
   ],
   "metadata": {}
  }
 ]
}